from pylab import *
import numpy

# 2D shapes

# quad
def f(xi, eta):
    return (1-xi**2)*(1-eta**2)

# triangle
def f(xi, eta):
    v = 27*xi*eta*(1.-xi-eta)
    minv = -0.1
    maxv = 1.1
    return v

f = numpy.vectorize(f)

X = linspace(-1,1,40)
Y = linspace(-1,1,40)
XX, YY = meshgrid(X,Y)

C = f(XX, YY)

p = contour(XX, YY, C, linspace(-0.1, 1.1, 12) )
colorbar(p)
show()

# 3D shapes

#from mpl_toolkits.mplot3d import
from mpl_toolkits.mplot3d import axes3d # for: projection='3d'
from matplotlib import pyplot

# hexa
def f(xi, eta, zeta):
    return (1-xi**2)*(1-eta**2)*(1-zeta**2)

# tet
def f(xi, eta, zeta):
    return 256*xi*eta*zeta*(1-xi-eta-zeta)

f = numpy.vectorize(f)

X = linspace(-1,1,40)
Y = linspace(-1,1,40)
Z = linspace(-1,1,40)
XX, YY = meshgrid(X,Y)
C = f(XX, YY, 0.5)

XX, YY = meshgrid(X,Y)

C = f(XX, YY, 0.25)

print C.size
p = contourf(XX, YY, C, linspace(-0.1, 1.01, 12) )
colorbar(p)
show()

